#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import pickle as pk
import cv2
import numpy as np
from pathlib import Path
from tqdm import tqdm

dataset_root="../datasets/icg"
annotation_path = os.path.join(dataset_root, "training_set", "gt", "bounding_box", "bounding_boxes", "person", "imgIdToBBoxArray.p")
image_path = os.path.join(dataset_root, "training_set", "images")

output_image_path = os.path.join(dataset_root, "training_set_1500x1000", "images")
output_annotation_path = os.path.join(dataset_root, "training_set_1500x1000", "gt", "bounding_box", "bounding_boxes", "person")

Path(output_image_path).mkdir(parents=True, exist_ok=True)
Path(output_annotation_path).mkdir(parents=True, exist_ok=True)

with open(annotation_path, "rb") as pickleFile:
    annotation_dict = pk.load(pickleFile)
    
new_annotation_dict = annotation_dict.copy()    
    
for image_id in tqdm(annotation_dict):
    img = cv2.imread(os.path.join(image_path, image_id + ".jpg"))
    img_resized = cv2.resize(img, (1500, 1000))
    cv2.imwrite(os.path.join(output_image_path, image_id + ".jpg"), img_resized)
    
    box_list = []
    for box in annotation_dict[image_id]:
        box_list.append(np.round(box/4))
        
    new_annotation_dict[image_id] = box_list
        
with open(os.path.join(output_annotation_path, "imgIdToBBoxArray.p"), 'wb') as pickleFile:
    pk.dump(new_annotation_dict, pickleFile, protocol=pk.HIGHEST_PROTOCOL)